import keras
from keras import layers
from tensorflow import saved_model
import numpy as np
import matplotlib.pyplot as plt
import os

number = 100

x1 = np.linspace(0, 0.5, number)
np.random.shuffle(x1)
y1 = np.linspace(0.5, 1.0, number)
np.random.shuffle(y1)
# 0 : 红色区域
z1 = np.zeros(x1.shape)

x1_tmp = np.array(x1, ndmin=2)
y1_tmp = np.array(y1, ndmin=2)
x1y1 = np.stack((x1_tmp, y1_tmp), axis=2)[0]

# 1 ：蓝色区域
x2 = np.linspace(0.5, 1, number)
np.random.shuffle(x2)
y2 = np.linspace(0, 0.5, number)
np.random.shuffle(y2)
z2 = np.ones(x2.shape)

x2_tmp = np.array(x2, ndmin=2)
y2_tmp = np.array(y2, ndmin=2)
x2y2 = np.stack((x2_tmp, y2_tmp), axis=2)[0]

plt.plot(x1, y1, 'ro')  # 红色点区域
plt.plot(x2, y2, 'bo')  # 蓝色点区域

# ---------------- 数据集 ---------------------

train_data = np.concatenate((x1y1, x2y2), axis=0)
train_labels = np.concatenate((z1, z2), axis=0).reshape(-1, 1)
# 将分类信息标注出来
train_labels = keras.utils.to_categorical(train_labels, num_classes=2)

# ----------------- 构建神经网络做分类训练 -------------

model = keras.Sequential([
    layers.Flatten(input_shape=(2,)),
    layers.Dense(10, activation="tanh"),
    layers.Dense(10, activation="tanh"),
    layers.Dense(2, activation="softmax")  # 结果是一个二分类
])

model.compile(
    # SGD 优化器 对于分类问题，梯度下降慢，需要增大学习率
    optimizer=keras.optimizers.SGD(learning_rate=0.01),
    loss=keras.losses.MSE,  # 损失率
    # accuracy 准确率
    metrics=["acc"]
)

# 训练时存在，训练集（平时考试），验证集（期末考试）
# model.fit(
#     # train_data 需要训练的数据
#     # train_labels 左边数据的分类
#     train_data, train_labels,
#     epochs=200,
#     validation_split=0.2)  # validation_split 从train_data分出百分之20的数据作为验证集
#
# # 将训练结果保存起来 注意保存的模型的格式是“.keras”
# model.save('model.keras')
#
# predict_labels = model.predict(train_data)
#
# # print(predict_labels)
# # 通过argmax 将一行中的最大概率的分类的下标
# show = np.argmax(predict_labels, axis=1)
# print(show)

plt.show()
